fix(speculative): fail closed when the streaming /done check cannot be reached - #2756

Failing closed on an unreachable /done is a correctness fix unrelated to training speed, so
it gets its own PR, #2756, which also narrows the except to the module's transient errors.
This restores hf_streaming_dataset.py and its test to main's.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

### What does this PR do?

Type of change: new feature (performance), bug fix

Speeds up DFlash-family drafter training. On Gemma-4-E4B DSpark
streaming (PDX B300, 3 nodes = 1 vLLM serve + 2 trainers, 16 DP ranks,
batch 4 × 4096), step time goes **1.690 → 0.484 s (3.49×)**, with the
loss unchanged within run-to-run noise. Every round below was measured
on the same window (steps 21-40, profiler off).

| Change | Step time |
|---|---|
| FlexAttention over a block-sparse draft mask
(`dflash_use_flex_attention`, opt-in). The dense `[B, 1, Q, KV]` mask
had forced SDPA onto its sm80-only memory-efficient kernel. | 1.690 →
0.711 s |
| Offline/streaming and DSpark online: project only the base-teacher
rows a loss reads, not full `[B, seq, vocab]` logits; one device sync
for DSpark's metrics instead of five | 0.711 → 0.629 s |
| DSpark TVD chunked with `Tensor.split` instead of slicing, whose
backward zero-filled a full `[N, vocab]` tensor per chunk | 0.629 →
0.541 s |
| `torch.compile` of the draft stack and the DSpark TVD chunk
(`dflash_use_torch_compile`, opt-in). With compile or FlexAttention on,
the anchor count is fixed by the config and `seq_len`, so the compiled
shapes stay static: recompiles in 40 steps drop from 4 to 1. Eager runs
still size it to the batch. | 0.541 → 0.484 s |

The same rounds, by phase and by GPU kernel:

<img width="2040" height="2040" alt="image"
src="https://github.com/user-attachments/assets/13ae02e5-a171-4f48-8739-c6828d018413"
/>

Also:
- Offline/streaming LiLiCorr (#2589) reads its candidate target logits
the same per-row way, instead of rebuilding full-sequence logits.
- FlexAttention's conflicts with the attention sink and with dropout are
rejected at convert time.
- No-behavior-change cleanups: the attention visibility rule is written
once (`_draft_mask_mod`) for SDPA and FlexAttention, and KD is switched
by `dflash_self_logit_distillation` rather than by a nullable argument.

**Why not HF's `flex_attention` dispatch (the path EAGLE takes)?** The
BlockMask is built here either way; the difference is the kernel call.
HF's `flex_attention_forward` passes `enable_gqa=True` whenever the
query-head count is a power of two, and that backward took 568 ms on the
Gemma-4-E4B drafter: 10× the path that repeats K/V first, and 2.5×
slower than the SDPA baseline. At `head_dim` > 256, FlexAttention's
default tiles fail on sm103 (misaligned address, launch failure), so the
kernel tiles are pinned to 32×32 and the BlockMask block size has to
divide them; `dflash_flex_attention.py` keeps both rules together.

### Usage

```yaml
dflash:
  dflash_use_flex_attention: true
  dflash_use_torch_compile: true
```

On multi-node runs, point `TRITON_CACHE_DIR` and
`TORCHINDUCTOR_CACHE_DIR` at node-local paths: on a shared filesystem,
concurrent ranks race on the compile cache.

### Testing

- `tests/unit/torch/speculative`: 387 passed on transformers 5.12.
`tests/gpu/torch/speculative/plugins`: 14 passed, 1 skipped.
- FlexAttention on torch 2.8.0 (manual, RTX PRO 6000): against an fp64
dense reference on padded-row shapes (seq 128, Q_LEN 144 and 244),
forward and backward agree to ≤ 7e-7 relative in fp32, and
compute-sanitizer memcheck is clean. Before the `q_block` clamp,
memcheck caught the backward reading past `anchors`.
- New tests cover FlexAttention vs. dense (GPU), the mask rule against a
hand-written mask and on padded FlexAttention rows, TVD chunk-size
invariance and `split` in the backward graph, anchor sampling vs. the
legacy sampler in eager and static modes, the per-row teacher (lm_head
never sees the full sequence offline), and FlexAttention conflicts at
convert time and through a `from_pretrained` resume.
- The cleanups were checked old vs. new: losses, gradients, dense masks
and GPU BlockMasks are bitwise equal, except DSpark's rnn head, which
differs at bf16 rounding.
- Against main, offline DFlash KD, LiLiCorr and DSpark get bitwise-equal
losses and gradients in bf16 from the per-row teacher. The DFlash
regression recipe (eager) trains in the same time as on main with an
identical loss curve.
- Step times are medians over 16 ranks × steps 21-40 (profiler off) on
the setup above.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ Both speedups are opt-in:
`dflash_use_flex_attention` and `dflash_use_torch_compile` default off.
The latter defaulted to True on main but did nothing; the shipped
recipes set it false.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
❌
- Did you get Claude approval on this PR?: ❌

### Additional Information

Based on main. The step times were measured on Gemma-4-E4B, which needs
#2186 to run; nothing in this PR is Gemma-specific.

The streaming `/done` fix that was part of this PR is now #2756.
Follow-up refactors are in draft #2757, stacked on this PR.

🤖 Generated with [Claude Code](https://claude.com/claude-code)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added an optional FlexAttention path for DFlash draft attention, with
block-sparse masking.
* Added optional training-time compilation for draft model components;
compilation is disabled by default.
* Teacher logits can now be obtained from supplied logits or captured
hidden states during distillation.
* **Bug Fixes**
* Improved draft training behavior for batches with no valid anchors,
ensuring trainable parameters receive gradients.
* Added validation for FlexAttention configurations that do not support
attention sinks or dropout.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Hao Guo <haoguo@nvidia.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>